import json
import os
import random
import re
import tqdm
import numpy as np
from multiprocessing import Pool
from functools import partial
import config
from api.gpt3_api import GPT3
from api.gpt4_api import GPT4
from api.kimi_api import KIMI
from api.llama3_api import LLAMA3
from api.global_config import GPT3_CONFIG, GPT4_CONFIG, KIMI_CONFIG, LLAMA3_CONFIG
from util.post_processing import calculate_rouge_scores


def sample_machine_instructions(machine_instructions, similarities, n):
    """Sample n machine instructions from a list of machine instructions."""
    return random.sample(machine_instructions, min(n, len(machine_instructions)))


def encode_prompt(prompt_instructions, is_chinese=False, handwrite_prompt=None):
    """Encode multiple prompt instructions into a single string."""
    if is_chinese:
        prompt = config.INSTRUCTION_PROMPT_CHINESE
    else:
        prompt = config.INSTRUCTION_PROMPT
    if handwrite_prompt != "":
        prompt = handwrite_prompt
    for idx, instruction in enumerate(prompt_instructions):
        # 去除空白字符且不以冒号结尾
        instruction = re.sub(r"\s+", " ", instruction).strip().rstrip(":")
        # 写作 1.xxx 2.xxx 3.xxx的形式
        prompt += f"{idx+1}. {instruction}\n"
    # 再多写一个序号，意思是让GPT继续续写
    prompt += f"{len(prompt_instructions) + 1}."
    return prompt


def get_model_instance(model_name):
    # 初始化每个config类
    gpt3_config = GPT3_CONFIG()
    gpt4_config = GPT4_CONFIG()
    # 现在使用kimi的补全功能，所以prompt可以调整一下
    completion_system_prompt = "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长句子补全，你需要根据提示补全用户的输入使之连贯完整。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力等问题的回答。Moonshot AI 为专有名词，不可翻译成其他语言。"
    kimi_config = KIMI_CONFIG(system_prompt=completion_system_prompt)
    llama3_config = LLAMA3_CONFIG()

    # 使用配置对象创建类的实例
    gpt3_instance = GPT3(gpt3_config)
    gpt4_instance = GPT4(gpt4_config)
    kimi_instance = KIMI(kimi_config)
    llama3_instance = LLAMA3(llama3_config)

    model_name_to_instances = {
        'gpt3': gpt3_instance,
        'gpt4': gpt4_instance,
        'kimi': kimi_instance,
        'llama3': llama3_instance,
    }
    if model_name in model_name_to_instances:
        # 获取对应的实例
        instance = model_name_to_instances[model_name]
    else:
        raise ValueError(f"No model found with the name '{model_name}'")
    return instance


if __name__ == "__main__":
    model_name = config.MODEL_NAME
    num_of_instructions = config.NUM_OF_INSTRUCTIONS
    num_prompt_instructions = config.NUM_PROMPT_INSTRUCTIONS
    handwrite_prompt = config.HANDWRITE_PROMPT

    is_chinese = False
    # batch_dir：不同模型生成数据的目录
    batch_dir = '../mydata/'+model_name+'_generations_'+config.BATCH_DIR_SUFFIX+'/'
    # seed_tasks_path：初始种子路径
    seed_tasks_path = config.SEED_TASKS_PATH
    # 调整is_chinese，因为有些模型的prompt更适合用中文来做，而其他模型英文prompt效果更好
    if model_name == "kimi":
        is_chinese = True


    seed_tasks = [json.loads(l) for l in open(seed_tasks_path, "r", encoding='utf-8')]
    seed_instructions = [item["instruction"] for item in seed_tasks]
    print(f"Loaded {len(seed_instructions)} human-written seed instructions")

    os.makedirs(batch_dir, exist_ok=True)
    request_idx = 0
    # 加载LLM生成的一部分指令，所以如果生成工作中断了也能继续做
    machine_instructions = []
    if os.path.exists(os.path.join(batch_dir, "machine_generated_instructions.jsonl")):
        with open(os.path.join(batch_dir, "machine_generated_instructions.jsonl"), "r", encoding='utf-8') as fin:
            for line in fin:
                instruction_info = json.loads(line)
                machine_instructions.append(instruction_info["instruction"])
                request_idx = instruction_info["request_idx"] + 1
        print(f"Loaded {len(machine_instructions)} machine-generated instructions")

    # now let's generate new instructions!
    progress_bar = tqdm.tqdm(total=num_of_instructions)
    if machine_instructions:
        progress_bar.update(len(machine_instructions))
    # 获取调用的LLM
    model_instance = get_model_instance(model_name)

    with open(os.path.join(batch_dir, "machine_generated_instructions.jsonl"), "a", encoding='utf-8') as fout:
        while len(machine_instructions) < num_of_instructions:
            # 选择2条机器生成的指令，如果一开始没有，那么列表为空
            prompt_instructions = sample_machine_instructions(
                machine_instructions,
                similarities=None,
                n=2)
            # sample human instructions from the pool，默认总数是8，前面2条这里就加6条，否则加8条
            prompt_instructions += random.sample(seed_instructions, num_prompt_instructions - len(prompt_instructions))
            random.shuffle(prompt_instructions)
            # 到现在为止，把机器生成的和人工编写的都混在一起了
            prompt = encode_prompt(prompt_instructions, is_chinese, handwrite_prompt)

            # 这个_必须加，因为要返回两个参数，否则格式不对
            instructions, _ = model_instance.get_instructions(prompt)

            print("[instructions]"+str(instructions))
            for inst in instructions:
                # 关于这一块的注释，详见test/test_rouge.py
                with Pool(4) as p:
                    rouge_scores = p.map(partial(calculate_rouge_scores, inst=inst), seed_instructions + machine_instructions)

                if max(rouge_scores) > 0.7:
                    print("["+inst + "]与当前指令池中指令相似度过高，舍去(max similiarity score=" + str(max(rouge_scores))+')')
                    continue

                all_instructions = seed_instructions + machine_instructions
                most_similar_instructions_indices = np.argsort(rouge_scores)[-10:][::-1]
                most_similar_instructions = {all_instructions[i]: rouge_scores[i] for i in
                                             most_similar_instructions_indices}

                machine_instructions.append(inst)
                fout.write(json.dumps({
                    "instruction": inst,
                    "most_similar": most_similar_instructions,
                    "avg_similarity_score": float(np.mean(rouge_scores)),
                    "request_idx": request_idx
                }) + "\n")

                # 每生成100条，抽5条打印，看看效果
                if progress_bar.n % 100 == 0:
                    sample_instructions = sample_machine_instructions(
                        machine_instructions,
                        similarities=None,
                        n=5)
                    print(sample_instructions)

                progress_bar.update(1)
            request_idx += 1
